import asyncio
import logging
import os
import re
import random
import shutil
import threading
import time
from datetime import datetime
from logging.handlers import RotatingFileHandler

from pydantic import json
from wxauto import WeChat
from openai import OpenAI


# ======================
# 配置中心
# ======================
class Config:
    # 监听配置（支持多账号多群聊）
    # 格式：[账号ID, 群聊名称, 提示词名称]
    LISTEN_LIST = [
        ["1", "AI测试1", "提示词示例2-1"],
        ["1", "AI测试", "提示词示例2-1"],  # 同一账号可监听多个群聊
        ["2", "AI测试", "提示词示例2-2"],  # 不同账号监听不同群聊
    ]

    # AI配置
    DEEPSEEK_API_KEY = "your-api-key"  # 替换为实际API Key
    DEEPSEEK_BASE_URL = "https://ark.cn-beijing.volces.com/api/v3"
    MODEL = "doubao-1-5-pro-256k-250115"
    MAX_TOKEN = 2000
    TEMPERATURE = 1.1

    # 记忆配置
    ENABLE_MEMORY = True
    MEMORY_DIR = "Memory_Temp"
    MAX_MESSAGE_LOG_ENTRIES = 4  # 触发记忆总结的消息数量
    MAX_MEMORY_NUMBER = 10  # 最大记忆存储数量

    # 消息处理配置
    QUIET_TIME = {  # 免打扰时间段（24小时制）
        "start": "22:00",
        "end": "08:00"
    }
    ACCEPT_ALL_GROUP_MESSAGES = False  # 是否接收所有群聊消息（否则仅处理@消息）


# ======================
# 日志系统
# ======================
logging.basicConfig(
    level=logging.INFO,
    format="%(asctime)s - %(levelname)s - %(message)s",
    handlers=[
        RotatingFileHandler("bot.log", maxBytes=10 * 1024 * 1024, backupCount=5),
        logging.StreamHandler()
    ]
)
logger = logging.getLogger(__name__)


# ======================
# 微信账户管理
# ======================
class WeChatAccount:
    """单个微信账户实例，支持多群聊监听"""

    def __init__(self, account_id):
        self.account_id = account_id
        self.wx = WeChat(hwnd=int(account_id))  # 微信窗口句柄
        self.robot_name = self.wx.A_MyIcon.Name  # 机器人在微信中的显示名称
        self.monitored_groups = self._init_monitored_groups()  # 初始化监听的群聊

    def _init_monitored_groups(self):
        """初始化监听的群聊列表"""
        groups = []
        for item in Config.LISTEN_LIST:
            if item[0] == self.account_id:
                groups.append(item[1])
                try:
                    self.wx.AddListenChat(who=item[1], savepic=True)
                    logger.info(f"账号{self.account_id}成功监听群聊：{item[1]}")
                except Exception as e:
                    logger.error(f"账号{self.account_id}监听群聊{item[1]}失败: {e}")
        return groups

    def send_message(self, group_name, content):
        """向指定群聊发送消息"""
        self.wx.ChatWith(group_name)  # 激活群聊窗口
        for part in content.split("\n"):
            if part:
                self.wx.SendMsg(part)
                # 模拟人类打字速度
                time.sleep(len(part) * random.uniform(0.01, 0.02))


class WeChatManager:
    """微信账户管理器，支持多账号管理"""

    _instances = {}  # 存储已创建的微信账户实例
    _lock = threading.Lock()  # 线程安全锁

    @classmethod
    def get_account(cls, account_id):
        """获取或创建微信账户实例（单例模式）"""
        with cls._lock:
            if account_id not in cls._instances:
                cls._instances[account_id] = WeChatAccount(account_id)
            return cls._instances[account_id]


# ======================
# 记忆管理系统
# ======================
class MemorySystem:
    """管理各群聊的记忆和会话历史"""

    def __init__(self, root_dir):
        self.root_dir = root_dir
        self.memory_dir = os.path.join(root_dir, Config.MEMORY_DIR)
        os.makedirs(self.memory_dir, exist_ok=True)

    def get_memory_path(self, account_id, group_name):
        """获取群聊记忆文件路径"""
        safe_name = re.sub(r'[\\/:*?"<>|]', '_', group_name)  # 替换非法文件名字符
        return os.path.join(self.memory_dir, f"{account_id}_{safe_name}_memories.json")

    def load_memories(self, account_id, group_name):
        """加载群聊历史记忆"""
        path = self.get_memory_path(account_id, group_name)
        return json.load(open(path, "r", encoding="utf-8")) if os.path.exists(path) else []

    def save_memories(self, account_id, group_name, memories):
        """保存群聊记忆"""
        path = self.get_memory_path(account_id, group_name)
        json.dump(memories, open(path, "w", encoding="utf-8"), ensure_ascii=False, indent=2)

    def log_message(self, account_id, group_name, content):
        """记录消息到临时日志"""
        log_file = os.path.join(self.memory_dir, f"{account_id}_{group_name}_chatlog.txt")
        with open(log_file, "a", encoding="utf-8") as f:
            f.write(f"[{datetime.now()}] {content}\n")
        self._clean_logs(log_file)  # 自动清理大文件

    def _clean_logs(self, path):
        """清理过大的日志文件"""
        if os.path.getsize(path) > 1024 * 1024:
            shutil.move(path, f"{path}.old_{int(time.time())}")


# ======================
# 核心聊天机器人
# ======================
class ChatBot:
    """微信聊天机器人核心类，处理消息和AI交互"""

    def __init__(self):
        self.root_dir = os.path.dirname(__file__)
        self.memory = MemorySystem(self.root_dir)
        self.client = OpenAI(
            api_key=Config.DEEPSEEK_API_KEY,
            base_url=Config.DEEPSEEK_BASE_URL
        )
        self.group_states = {}  # 存储各群聊状态

    def _get_group_state(self, account_id, group_name):
        """获取群聊会话状态（懒加载）"""
        key = f"{account_id}_{group_name}"
        if key not in self.group_states:
            self.group_states[key] = {
                "messages": [],  # 当前会话消息队列
                "memories": self.memory.load_memories(account_id, group_name)  # 历史记忆
            }
        return self.group_states[key]

    def _is_quiet_time(self):
        """判断是否处于免打扰时间"""
        now = datetime.now().time()
        start = datetime.strptime(Config.QUIET_TIME["start"], "%H:%M").time()
        end = datetime.strptime(Config.QUIET_TIME["end"], "%H:%M").time()
        return start <= now <= end if start <= end else now >= start or now <= end

    async def process_group_message(self, msg, account_id, group_name):
        """处理群聊消息"""
        state = self._get_group_state(account_id, group_name)
        content = self._clean_message(msg, account_id, group_name)

        if not content:
            return

        logger.info(f"[账号:{account_id}/群聊:{group_name}] 收到消息: {content}")
        self.memory.log_message(account_id, group_name, f"用户: {content}")
        state["messages"].append(content)

        # 消息数量超过限制时清理旧消息
        if len(state["messages"]) > 10:
            state["messages"].pop(0)

        await self._generate_response(account_id, group_name)

    def _clean_message(self, msg, account_id, group_name):
        """清理消息内容（处理@和无效消息）"""
        content = msg.content if hasattr(msg, "content") else ""

        # 处理群聊消息
        if msg.type == "group":
            robot_name = WeChatManager.get_account(account_id).robot_name

            if Config.ACCEPT_ALL_GROUP_MESSAGES:
                return f"[群聊@{msg.sender}] {content}"
            else:
                # 仅处理@机器人的消息
                if f"@{robot_name}" in content:
                    # 移除@标记并清理内容
                    cleaned_content = re.sub(f"@{robot_name}[\u2005\\s]*", "", content).strip()
                    return f"[群聊@{msg.sender}] {cleaned_content}"
                return None  # 忽略未@的群聊消息

        return None  # 忽略非群聊消息

    async def _generate_response(self, account_id, group_name):
        """生成AI回复"""
        state = self._get_group_state(account_id, group_name)
        prompt = self._build_prompt(account_id, group_name, state["messages"])

        try:
            response = self.client.chat.completions.create(
                model=Config.MODEL,
                messages=prompt,
                temperature=Config.TEMPERATURE,
                max_tokens=Config.MAX_TOKEN
            )
            reply = response.choices[0].message.content.strip()
            await self._send_reply(account_id, group_name, reply)
            self._update_memory(account_id, group_name, state["messages"], reply)
            state["messages"] = []  # 清空当前消息队列

        except Exception as e:
            logger.error(f"生成回复失败: {e}")
            await self._send_reply(account_id, group_name, "抱歉，暂时无法回复")

    def _build_prompt(self, account_id, group_name, messages):
        """构建包含历史记忆的提示词"""
        state = self._get_group_state(account_id, group_name)

        # 获取群聊对应的提示词配置
        prompt_name = next((item[2] for item in Config.LISTEN_LIST
                            if item[0] == account_id and item[1] == group_name), "default")

        # 加载系统提示词
        system_prompt = self._load_system_prompt(prompt_name)

        # 构建历史记忆上下文
        memory_context = "\n".join([f"历史记忆: {mem}" for mem in state["memories"][-Config.MAX_MEMORY_NUMBER:]])

        # 构建用户消息上下文
        user_context = "\n".join(messages)

        return [
            {"role": "system", "content": system_prompt},
            {"role": "system", "content": memory_context},
            {"role": "user", "content": user_context}
        ]

    def _load_system_prompt(self, prompt_name):
        """加载系统提示词文件"""
        prompt_path = os.path.join(self.root_dir, "prompts", f"{prompt_name}.md")
        if os.path.exists(prompt_path):
            return open(prompt_path, "r", encoding="utf-8").read()
        return "你是一个智能助手，需根据历史对话和当前问题回答用户"

    def _update_memory(self, account_id, group_name, messages, reply):
        """更新群聊记忆"""
        if not Config.ENABLE_MEMORY:
            return

        # 生成对话摘要
        conversation = "\n".join([f"用户: {msg}" for msg in messages] + [f"机器人: {reply}"])
        summary = self._summarize_conversation(conversation)

        if summary:
            state = self._get_group_state(account_id, group_name)
            state["memories"].append(summary)

            # 限制记忆数量
            if len(state["memories"]) > Config.MAX_MEMORY_NUMBER:
                state["memories"] = state["memories"][-Config.MAX_MEMORY_NUMBER:]

            self.memory.save_memories(account_id, group_name, state["memories"])

    def _summarize_conversation(self, conversation):
        """调用AI生成对话摘要"""
        prompt = f"请总结以下对话（不超过200字）：\n{conversation}"
        try:
            response = self.client.chat.completions.create(
                model=Config.MODEL,
                messages=[{"role": "user", "content": prompt}],
                temperature=0.5,
                max_tokens=200
            )
            return response.choices[0].message.content.strip()
        except Exception:
            return None

    async def _send_reply(self, account_id, group_name, content):
        """发送回复（处理免打扰逻辑）"""
        if self._is_quiet_time():
            logger.info("免打扰时间，不发送消息")
            return

        try:
            account = WeChatManager.get_account(account_id)
            account.send_message(group_name, content)
            logger.info(f"[账号:{account_id}/群聊:{group_name}] 回复成功: {content}")
        except Exception as e:
            logger.error(f"发送回复失败: {e}")

    async def start_listening(self):
        """启动多账号多群聊监听"""
        logger.info(f"启动监听 {len(Config.LISTEN_LIST)} 个群聊配置...")

        # 获取所有账号ID
        account_ids = set(item[0] for item in Config.LISTEN_LIST)

        # 为每个账号创建监听任务
        tasks = []
        for account_id in account_ids:
            tasks.append(self._listen_account(account_id))

        # 并发执行所有监听任务
        await asyncio.gather(*tasks)

    async def _listen_account(self, account_id):
        """监听单个账号的所有群聊"""
        account = WeChatManager.get_account(account_id)
        logger.info(f"账号{account_id}开始监听群聊")

        while True:
            try:
                # 获取该账号监听的所有群聊消息
                for group_name in account.monitored_groups:
                    msgs = account.wx.GetListenMessage(group_name)
                    for msg in msgs:
                        if msg.type == "group":
                            await self.process_group_message(msg, account_id, group_name)
            except Exception as e:
                logger.error(f"账号{account_id}监听异常: {e}")
            await asyncio.sleep(0.1)


# ======================
# 程序入口
# ======================
async def main():
    bot = ChatBot()
    try:
        await bot.start_listening()
    except KeyboardInterrupt:
        logger.info("用户终止程序")
    except Exception as e:
        logger.error(f"程序异常: {e}", exc_info=True)


if __name__ == "__main__":
    asyncio.run(main())